package com.aliyun.solutions.demo.wanxiang.web;

import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisOutput;
import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisParam;
import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisResult;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.aliyun.oss.OSS;
import com.aliyun.oss.OSSClientBuilder;
import com.aliyun.oss.common.auth.CredentialsProviderFactory;
import com.aliyun.oss.common.auth.EnvironmentVariableCredentialsProvider;
import com.aliyun.oss.model.PutObjectRequest;
import com.aliyun.solutions.demo.wanxiang.model.ApiResult;
import com.aliyun.solutions.demo.wanxiang.model.param.FaceChainGenerationParam;
import com.aliyun.solutions.demo.wanxiang.model.param.GraffitiPaintingParam;
import com.aliyun.solutions.demo.wanxiang.model.param.TextToImageParam;
import com.aliyuncs.exceptions.ClientException;
import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import okhttp3.Headers.Builder;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.Validate;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ClassPathResource;
import org.springframework.util.CollectionUtils;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.ResponseBody;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import javax.annotation.PostConstruct;
import javax.validation.Valid;
import javax.validation.constraints.NotEmpty;
import javax.validation.constraints.NotNull;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.time.Duration;
import java.util.Base64;
import java.util.Date;
import java.util.List;
import java.util.Properties;

/**
 * API Controller
 *
 * @author solution
 * @date 2024/12/04
 */
@Slf4j
@Validated
@RestController
@RequestMapping("/api")
public class ApiController {

    private static final OkHttpClient OK_HTTP_CLIENT = new OkHttpClient.Builder().callTimeout(Duration.ofSeconds(60))
        .connectTimeout(Duration.ofSeconds(10)).readTimeout(Duration.ofSeconds(10)).writeTimeout(Duration.ofSeconds(10))
        .build();

    private static OSS ossClient = null;

    private static final String TEXT_TO_IMAGE_API
        = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis";

    private static final String IMAGE_TO_IMAGE_API
        = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis/";

    private static final String IMAGE_GENERATION_API
        = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image-generation/generation";

    private static final String FACECHAIN_GENERATION_API
        = "https://dashscope.aliyuncs.com/api/v1/services/aigc/album/gen_potrait";

    private static final String GET_TASK_RESULT_API = "https://dashscope.aliyuncs.com/api/v1/tasks/%s";

    private static final Gson GSON = new Gson();
    private static final String JSON_CONTENT_TYPE = "application/json; charset=utf-8";
    private static final String AUTHORIZATION = "Authorization";
    private static final String ENABLE = "enable";
    private static final String X_DASH_SCOPE_ASYNC = "X-DashScope-Async";
    private static final String BEARER = "Bearer ";

    private static final String WANX_V1 = "wanx-v1";
    private static final String WANX_SKETCH_TO_IMAGE_LITE = "wanx-sketch-to-image-lite";
    private static final String WANX_STYLE_REPAINT_V_1 = "wanx-style-repaint-v1";
    private static final String FACECHAIN_GENERATION = "facechain-generation";

    private static final String DATA_IMAGE_PNG_BASE_64 = "data:image/png;base64,";
    private static final String PNG = "png";

    private static final String DASHSCOPE_API_KEY = "DASHSCOPE_API_KEY";
    private static final String OSS_BUCKET = "OSS_BUCKET";
    private static final String OSS_ENDPOINT = "OSS_ENDPOINT";

    private static final String CODE = "code";
    private static final String MESSAGE = "message";



    public static final Properties ERROR_MESSAGES = new Properties();

    static {
        try (InputStream inputStream = new ClassPathResource("/error_message.properties").getInputStream()) {
            ERROR_MESSAGES.load(inputStream);
        } catch (IOException e) {
            throw new RuntimeException("Failed to load error_message.properties", e);
        }
    }

    @PostConstruct
    private void init() {
        // 从环境变量中获取访问凭证。运行本代码示例之前，请确保已设置环境变量 OSS_ACCESS_KEY_ID 和 OSS_ACCESS_KEY_SECRET。
        EnvironmentVariableCredentialsProvider credentialsProvider = null;
        try {
            credentialsProvider = CredentialsProviderFactory.newEnvironmentVariableCredentialsProvider();
        } catch (ClientException e) {
            throw new RuntimeException("服务器上未设置环境变量：OSS_ACCESS_KEY_ID、OSS_ACCESS_KEY_SECRET");
        }

        Validate.notEmpty(getDashScopeApiKey(), "服务器上未设置环境变量：%s", DASHSCOPE_API_KEY);
        Validate.notEmpty(getOssEndpoint(), "服务器上未设置环境变量：%s", OSS_ENDPOINT);
        Validate.notEmpty(getOssBucketName(), "服务器上未设置环境变量：%s", OSS_BUCKET);

        // 创建OSSClient实例。
        ossClient = new OSSClientBuilder().build(getOssEndpoint(), credentialsProvider);
    }

    /**
     * 文生图
     *
     * @param parameter 参数
     * @return {@link ImageSynthesisResult }
     * @throws Exception 异常
     */
    @PostMapping("/text-to-image")
    @ResponseBody
    public ApiResult<ImageSynthesisOutput> textToImage(@RequestBody @NotNull @Valid TextToImageParam parameter)
        throws Exception {
        ImageSynthesisParam param = ImageSynthesisParam.builder().model(WANX_V1).n(parameter.getNum()).size(
            parameter.getSize()).prompt(parameter.getPrompt()).negativePrompt(parameter.getNegativePrompt()).style(
            parameter.getStyles()).build();
        // 请求ai
        return callApi(TEXT_TO_IMAGE_API, param.getHttpBody().toString());
    }

    /**
     * 涂鸦作画
     *
     * @param parameter 参数
     * @return {@link ImageSynthesisResult }
     * @throws Exception 异常
     */
    @PostMapping("/graffiti-painting")
    @ResponseBody
    public ApiResult<ImageSynthesisOutput> graffitiPainting(
        @RequestBody @NotNull @Valid GraffitiPaintingParam parameter) throws Exception {
        // 上传图片
        String imageData = parameter.getImageData().replace(DATA_IMAGE_PNG_BASE_64, "");
        byte[] bytes = Base64.getDecoder().decode(imageData);
        String url = uploadInputStreamToOss(new ByteArrayInputStream(bytes), PNG);

        // 请求ai
        ImageSynthesisParam param = ImageSynthesisParam.builder().model(WANX_SKETCH_TO_IMAGE_LITE).n(parameter.getNum())
            .size(parameter.getSize()).prompt(parameter.getPrompt()).sketchImageUrl(url).style(parameter.getStyles())
            .build();
        return callApi(IMAGE_TO_IMAGE_API, param.getHttpBody().toString());
    }

    /**
     * 人物写真生成
     *
     * @param param 参数
     * @return {@link ImageSynthesisResult }
     * @throws Exception 异常
     */
    @PostMapping("/facechain-generation")
    @ResponseBody
    public ApiResult<ImageSynthesisOutput> faceChainGeneration(@NotNull FaceChainGenerationParam param)
        throws Exception {
        log.info("param:{}", param);

        // 上传图片
        JsonArray userImageUrlList = new JsonArray();
        if (!CollectionUtils.isEmpty(param.getUserImageList())) {
            for (MultipartFile file : param.getUserImageList()) {
                if (file.getSize() == 0) {
                    continue;
                }
                String suffix = file.getOriginalFilename().substring(file.getOriginalFilename().lastIndexOf(".") + 1);
                String url = uploadInputStreamToOss(file.getInputStream(), suffix);
                userImageUrlList.add(url);
            }
        }
        String templateUrl = null;
        if (param.getTemplateImage() != null && param.getTemplateImage().getSize() > 0) {
            MultipartFile file = param.getTemplateImage();
            String suffix = file.getOriginalFilename().substring(file.getOriginalFilename().lastIndexOf(".") + 1);
            templateUrl = uploadInputStreamToOss(file.getInputStream(), suffix);
        }

        // 请求ai
        JsonObject json = new JsonObject();
        json.addProperty("model", FACECHAIN_GENERATION);
        JsonObject parameters = new JsonObject();
        parameters.addProperty("style", param.getStyles());
        parameters.addProperty("size", param.getSize());
        parameters.addProperty("n", param.getNum());
        json.add("parameters", parameters);
        JsonObject input = new JsonObject();
        input.add("user_urls", userImageUrlList);
        input.addProperty("template_url", templateUrl);
        json.add("input", input);
        log.info(GSON.toJson(json));
        return callApi(FACECHAIN_GENERATION_API, GSON.toJson(json));
    }

    /**
     * 图像生成
     *
     * @param image      图像
     * @param styleIndex 风格指数
     * @return {@link ImageSynthesisResult }
     * @throws Exception 异常
     */
    @PostMapping("/image-generation")
    @ResponseBody
    public ApiResult<ImageSynthesisOutput> imageGeneration(@NotNull(message = "图片不能为空") MultipartFile image,
        @NotNull(message = "风格不能为空") Integer styleIndex) throws Exception {
        // 上传图片
        String suffix = image.getOriginalFilename().substring(image.getOriginalFilename().lastIndexOf(".") + 1);
        String url = uploadInputStreamToOss(image.getInputStream(), suffix);

        // 请求ai
        JsonObject json = new JsonObject();
        json.addProperty("model", WANX_STYLE_REPAINT_V_1);
        JsonObject input = new JsonObject();
        input.addProperty("image_url", url);
        input.addProperty("style_index", styleIndex);
        json.add("input", input);
        return callApi(IMAGE_GENERATION_API, GSON.toJson(json));
    }

    /**
     * 调用api
     *
     * @param imageGenerationApi 图像生成API
     * @param jsonStr            json str
     * @return {@link ApiResult }<{@link ImageSynthesisOutput }>
     * @throws Exception 异常
     */
    private ApiResult<ImageSynthesisOutput> callApi(String imageGenerationApi, String jsonStr) throws Exception {
        MediaType contentType = MediaType.get(JSON_CONTENT_TYPE);
        Headers headers = new Builder().add(AUTHORIZATION, BEARER + getDashScopeApiKey()).add(X_DASH_SCOPE_ASYNC,
            ENABLE).build();
        Request request = new Request.Builder().url(imageGenerationApi).headers(headers).post(
            okhttp3.RequestBody.create(contentType, jsonStr)).build();
        Response response = OK_HTTP_CLIENT.newCall(request).execute();
        return handleResponse(response);
    }

    /**
     * 获取任务结果
     *
     * @param taskId 任务id
     * @return {@link ImageSynthesisResult }
     * @throws Exception 异常
     */
    @GetMapping("/task-result")
    @ResponseBody
    public ApiResult<ImageSynthesisOutput> getTaskResult(@NotEmpty(message = "taskId不能为空") String taskId)
        throws Exception {
        Headers headers = new Builder().add(AUTHORIZATION, BEARER + getDashScopeApiKey()).build();
        String url = String.format(GET_TASK_RESULT_API, taskId);
        Request request = new Request.Builder().url(url).headers(headers).get().build();
        Response response = OK_HTTP_CLIENT.newCall(request).execute();
        return handleResponse(response);
    }

    /**
     * 处理响应
     *
     * @param response 响应
     * @return {@link ApiResult }<{@link ImageSynthesisOutput }>
     * @throws IOException ioexception
     */
    private ApiResult<ImageSynthesisOutput> handleResponse(Response response) throws IOException {
        // 检查响应是否成功
        okhttp3.ResponseBody body = response.body();
        String bodyString = body == null ? null : body.string();
        if (body == null || StringUtils.isBlank(bodyString)) {
            throw new RuntimeException("向AI发送请求失败，请重试！message=" + JSON.toJSONString(response));
        }
        log.info("httpResponseCode: {} httpResponseBody: {}", response.code(), bodyString);
        JSONObject jsonObject = JSON.parseObject(bodyString);
        String errorCode = jsonObject.getString(CODE);
        String errorMsg = jsonObject.getString(MESSAGE);
        if (jsonObject.containsKey(CODE)) {
            checkResult(errorCode, errorMsg, response.code());
        }
        ImageSynthesisResult result = GSON.fromJson(bodyString, ImageSynthesisResult.class);
        errorCode = result.getOutput().getCode();
        errorMsg = result.getOutput().getMessage();
        checkResult(errorCode, errorMsg, response.code());
        return ApiResult.success(result.getOutput());
    }

    /**
     * 检查结果
     *
     * @param errorCode 错误代码
     * @param errorMsg  错误味精
     * @param httpCode  http代码
     */
    private void checkResult(String errorCode, String errorMsg, int httpCode) {
        if (StringUtils.isBlank(errorCode)) {
            return;
        }
        String message = ERROR_MESSAGES.getProperty(httpCode + "_" + errorCode);
        if (StringUtils.isBlank(message)) {
            message = ERROR_MESSAGES.getProperty("*_" + errorCode);
        }
        if (StringUtils.isNotBlank(message)) {
            if (message.startsWith("[")) {
                List<String> messageArray = JSON.parseArray(message, String.class);
                for (int i = 0; i <= messageArray.size() / 2; i += 2) {
                    if (errorMsg.contains(messageArray.get(i))) {
                        message = messageArray.get(i + 1);
                        break;
                    }
                }
            }
        }
        if (StringUtils.isBlank(message)) {
            String targetKey = (String)ERROR_MESSAGES.keySet().stream().filter(
                p -> StringUtils.substringAfter(p.toString(), "_").equals(errorCode)).findFirst().orElse(null);
            if (StringUtils.isNotBlank(targetKey)) {
                message = ERROR_MESSAGES.getProperty(targetKey);
            } else {
                message = errorMsg;
            }
        }
        throw new RuntimeException("向AI发送请求失败，请重试！message=" + message);
    }

    /**
     * 上传输入流到oss
     *
     * @param inputStream 输入流
     * @param suffix      后缀
     * @return {@link String }
     */
    private String uploadInputStreamToOss(InputStream inputStream, String suffix) {
        String osspath = String.format("ai_test_images/%d_%s.%s", System.currentTimeMillis(),
            RandomStringUtils.random(8, true, true), suffix);

        try {
            // 上传
            String ossBucketName = getOssBucketName();
            PutObjectRequest putObjectRequest = new PutObjectRequest(ossBucketName, osspath, inputStream);
            ossClient.putObject(putObjectRequest);
            // 设置签名URL过期时间，单位为毫秒。本示例以设置过期时间为1小时为例。
            Date expiration = new Date(System.currentTimeMillis() + 3600 * 1000L);
            // 生成以GET方法访问的签名URL，访客可以直接通过浏览器访问相关内容。
            URL url = ossClient.generatePresignedUrl(ossBucketName, osspath, expiration);
            return url.toString();
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            throw new RuntimeException("上传oss失败, message=" + e.getMessage());
        }
    }

    /**
     * 获取 API密钥
     *
     * @return {@link String }
     */
    private String getDashScopeApiKey() {
        return StringUtils.trim(System.getenv(DASHSCOPE_API_KEY));
    }

    private String getOssBucketName() {
        return StringUtils.trim(System.getenv(OSS_BUCKET));
    }

    private String getOssEndpoint() {
        return StringUtils.trim(System.getenv(OSS_ENDPOINT));
    }
}